#!/bin/bash
# ThingsBoard 增量还原脚本
# 版本: 3.0.0 (优化简化版 - 支持直接恢复模式)
#
# 优化说明:
# 本脚本采用模块化设计，分为以下几个主要部分：
# 1. 初始化 - 加载配置、设置环境变量
# 2. 辅助函数 - 通用功能函数
# 3. 恢复流程函数 - 分步骤执行恢复过程
# 4. 主程序入口 - 解析参数，调用相应的恢复流程

# ========== 1. 初始化部分 ==========

# 设置脚本名称
SCRIPT_NAME="restore_incremental"

# 记录当前工作目录
CURRENT_DIR=$(pwd)

# 获取脚本所在目录的绝对路径
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)

# 根据脚本目录计算项目根目录路径
if [[ "$SCRIPT_DIR" == */scripts/restore ]]; then
  # 如果脚本位于scripts/restore目录中
  PROJECT_ROOT=$(cd "$SCRIPT_DIR/../.." && pwd)
elif [[ "$SCRIPT_DIR" == */scripts/* ]]; then
  # 如果脚本位于scripts目录或其他子目录中
  PROJECT_ROOT=$(echo "$SCRIPT_DIR" | sed 's|\(.*\)/scripts.*|\1|')
else
  # 否则假设当前目录是项目根目录
  PROJECT_ROOT=$(pwd)
fi

# 配置文件路径
CONFIG_FILE="$PROJECT_ROOT/config/backup_env.conf"

# 检查配置文件是否存在
if [ ! -f "$CONFIG_FILE" ]; then
  echo "错误: 配置文件不存在: $CONFIG_FILE"
  exit 1
fi

# 导入共用函数 - 使用绝对路径
COMMON_SCRIPT="$PROJECT_ROOT/scripts/common.sh"
if [ -f "$COMMON_SCRIPT" ]; then
  source "$COMMON_SCRIPT"
else
  echo "错误: 共用脚本不存在: $COMMON_SCRIPT"
  exit 1
fi

# ========== 2. 辅助函数 ==========

# 创建WAL-G配置文件
create_walg_config() {
  local prefix=$1
  cat > "/root/.walg.json" <<EOF
{
  "WALG_S3_PREFIX": "s3://${MINIO_BUCKET}/${prefix}",
  "AWS_ACCESS_KEY_ID": "${MINIO_USER}",
  "AWS_SECRET_ACCESS_KEY": "${MINIO_PASSWORD}",
  "AWS_ENDPOINT": "http://${MINIO_HOST}:${MINIO_PORT}",
  "AWS_S3_FORCE_PATH_STYLE": "true",
  "PGHOST": "localhost",
  "PGPORT": "5432",
  "PGUSER": "$POSTGRES_SUPERUSER",
  "PGPASSWORD": "$POSTGRES_PASSWORD",
  "PGDATABASE": "$THINGSBOARD_DATABASE_NAME",
  "WALG_COMPRESSION_METHOD": "lz4"
}
EOF
  log_info "WAL-G配置文件已创建: /root/.walg.json (前缀: ${prefix})"
}

# 默认创建配置，使用增量备份前缀
create_walg_config "${MINIO_DIFF_PREFIX}/basebackups_005"

# 设置WAL-G环境变量
set_walg_env() {
  local prefix=$1
  export WALG_S3_PREFIX="s3://${MINIO_BUCKET}/${prefix}"
  export AWS_ACCESS_KEY_ID="${MINIO_USER}"
  export AWS_SECRET_ACCESS_KEY="${MINIO_PASSWORD}"
  export AWS_ENDPOINT="http://${MINIO_HOST}:${MINIO_PORT}"
  export AWS_S3_FORCE_PATH_STYLE="true"
  export PGHOST=localhost
  export PGPORT=5432
  export PGUSER=$POSTGRES_SUPERUSER
  export PGPASSWORD=$POSTGRES_PASSWORD
  export WALG_COMPRESSION_METHOD=lz4

  log_info "WAL-G环境变量已设置 (前缀: ${prefix})"
}

# 显示帮助信息
show_usage() {
    echo "使用方法: $0 [选项] <备份名称>"
    echo "选项:"
    echo "  -h, --help      显示帮助信息"
    echo "  -l, --list      列出可用的备份"
    echo "  -f, --full      强制使用全量备份模式"
    echo "  -d, --diff      强制使用增量备份模式"
    echo "  -D, --direct    使用直接恢复模式（在主机上执行WAL-G，而不是在Docker容器中）"
    echo "备份名称可以是完整的备份ID（如base_00000001000000000000000A）或关键字'latest'来还原最新备份"
    exit 1
}

# 列出所有可用备份
list_backups() {
    echo "可用的备份列表:"
    echo "================="

    # 设置WAL-G环境变量
    set_walg_env "${MINIO_FULL_PREFIX}/basebackups_005"

    # 检查全量备份
    echo -e "\n全量备份列表:"
    echo "================="
    FULL_BACKUP_LIST=$($(which wal-g) backup-list 2>&1)
    echo "$FULL_BACKUP_LIST"

    # 设置WAL-G环境变量为增量备份
    set_walg_env "${MINIO_DIFF_PREFIX}/basebackups_005"

    # 获取增量备份
    echo -e "\n增量备份列表:"
    echo "================="
    DIFF_BACKUP_LIST=$($(which wal-g) backup-list 2>&1)
    echo "$DIFF_BACKUP_LIST"

    # 获取增量备份详情
    echo -e "\n增量备份详情:"
    echo "================="
    for backup in $(echo "$DIFF_BACKUP_LIST" | grep base_ | awk '{print $1}'); do
        if [ -n "$backup" ]; then
            echo -e "\n备份: $backup"
            $(which wal-g) backup-fetch --detail "$backup" 2>&1
        fi
    done

    exit 0
}

# 检测备份位置
detect_backup_location() {
    log_info "自动检测备份位置..."

    # 检查全量备份
    set_walg_env "${MINIO_FULL_PREFIX}/basebackups_005"
    FULL_BACKUP_EXISTS=$($(which wal-g) backup-list 2>/dev/null | grep -c "$BACKUP_NAME" || echo "0")

    # 检查增量备份
    set_walg_env "${MINIO_DIFF_PREFIX}/basebackups_005"
    DIFF_BACKUP_EXISTS=$($(which wal-g) backup-list 2>/dev/null | grep -c "$BACKUP_NAME" || echo "0")

    if [ "$FULL_BACKUP_EXISTS" -eq "1" ]; then
        BACKUP_MODE="full"
        log_info "在全量备份中找到备份: $BACKUP_NAME"
    elif [ "$DIFF_BACKUP_EXISTS" -eq "1" ]; then
        BACKUP_MODE="diff"
        log_info "在增量备份中找到备份: $BACKUP_NAME"
    else
        log_error "指定的备份 $BACKUP_NAME 在全量和增量备份中均未找到"
        exit 1
    fi

    # 设置正确的S3前缀
    if [ "$BACKUP_MODE" = "full" ]; then
        set_walg_env "${MINIO_FULL_PREFIX}/basebackups_005"
        create_walg_config "${MINIO_FULL_PREFIX}/basebackups_005"
        log_info "使用全量备份模式，S3前缀: s3://${MINIO_BUCKET}/${MINIO_FULL_PREFIX}/basebackups_005"
    else
        set_walg_env "${MINIO_DIFF_PREFIX}/basebackups_005"
        create_walg_config "${MINIO_DIFF_PREFIX}/basebackups_005"
        log_info "使用增量备份模式，S3前缀: s3://${MINIO_BUCKET}/${MINIO_DIFF_PREFIX}/basebackups_005"
    fi
}

# 获取最新备份名称
get_latest_backup() {
    local step=$1
    log_info "获取最新备份..." "$step"

    BACKUP_NAME=$($(which wal-g) backup-list | tail -n 1 | awk '{print $1}')
    if [ -z "$BACKUP_NAME" ]; then
        handle_error 30 "未找到可用的备份" "$step"
    fi
    log_info "找到最新备份: $BACKUP_NAME" "$step"
}

# ========== 3. 恢复流程函数 ==========

# 停止服务
stop_services() {
    local step=$1
    log_info "准备停止服务..." "$step"

    # 停止ThingsBoard服务
    if [ "$(systemctl is-active thingsboard)" = "active" ]; then
        log_info "停止ThingsBoard服务..." "$step"
        systemctl stop thingsboard
        sleep 5
    fi

    # 停止PostgreSQL服务
    if [ "$USE_DOCKER" = "true" ]; then
        log_info "停止PostgreSQL容器..." "$step"
        docker stop "$POSTGRES_CONTAINER_NAME"
        sleep 5
    else
        log_info "停止PostgreSQL服务..." "$step"
        if systemctl is-active postgresql &>/dev/null; then
            systemctl stop postgresql
            sleep 5
        fi
    fi

    log_success "所有服务已停止" "$step"
}

# 准备数据目录
prepare_data_directory() {
    local step=$1
    log_info "准备PostgreSQL数据目录..." "$step"

    # 确保数据目录的父目录存在
    mkdir -p "$(dirname "$POSTGRES_DATA_DIR")"

    # 备份现有数据目录（如果存在）
    if [ -d "$POSTGRES_DATA_DIR" ]; then
        POSTGRES_DATA_BACKUP="$POSTGRES_DATA_DIR.bak.$(date +%Y%m%d%H%M%S)"
        log_info "备份现有数据目录到 $POSTGRES_DATA_BACKUP..." "$step"
        mv "$POSTGRES_DATA_DIR" "$POSTGRES_DATA_BACKUP"
    fi

    # 创建空的数据目录
    mkdir -p "$POSTGRES_DATA_DIR"
    log_success "PostgreSQL数据目录已准备就绪" "$step"
}

# 在容器中执行还原
restore_in_container() {
    local step=$1
    local backup_name=$2

    log_info "在Docker容器中执行WAL-G恢复..." "$step"

    # 创建一个临时容器来执行恢复
    local temp_container="temp-postgres-restore"
    log_info "创建临时PostgreSQL容器进行恢复..." "$step"

    # 停止原有容器（如果运行中）
    docker stop "$POSTGRES_CONTAINER_NAME" || log_warn "PostgreSQL容器已经停止" "$step"

    # 运行临时容器
    docker run -d --name "$temp_container" \
        -v "$POSTGRES_DATA_DIR:/var/lib/postgresql/data" \
        -v "/backup:/backup" \
        -e POSTGRES_PASSWORD="$POSTGRES_PASSWORD" \
        --network thingbelt-network \
        registry.cn-beijing.aliyuncs.com/xscha/base:postgres_14_4 sleep infinity

    # 检查容器是否创建成功
    if ! docker ps | grep -q "$temp_container"; then
        handle_error 35 "无法创建临时PostgreSQL容器" "$step"
    fi

    # 复制WAL-G到临时容器
    local walg_path=$(which wal-g)
    if [ -z "$walg_path" ]; then
        handle_error 36 "未找到wal-g命令，请先安装WAL-G" "$step"
    fi

    log_info "复制WAL-G到临时容器..." "$step"
    docker cp "$walg_path" "$temp_container:/usr/local/bin/"
    docker exec "$temp_container" chmod +x /usr/local/bin/wal-g

    # 在临时容器中执行恢复
    log_info "在临时容器中执行WAL-G恢复..." "$step"
    docker exec -e WALG_S3_PREFIX="$WALG_S3_PREFIX" \
                -e AWS_ACCESS_KEY_ID="$AWS_ACCESS_KEY_ID" \
                -e AWS_SECRET_ACCESS_KEY="$AWS_SECRET_ACCESS_KEY" \
                -e AWS_ENDPOINT="$AWS_ENDPOINT" \
                -e AWS_S3_FORCE_PATH_STYLE="$AWS_S3_FORCE_PATH_STYLE" \
                -e PGHOST="$PGHOST" \
                -e PGPORT="$PGPORT" \
                -e PGUSER="$PGUSER" \
                -e PGPASSWORD="$PGPASSWORD" \
                -e WALG_COMPRESSION_METHOD="$WALG_COMPRESSION_METHOD" \
                "$temp_container" bash -c "/usr/local/bin/wal-g backup-fetch /var/lib/postgresql/data $backup_name"

    RESTORE_STATUS=$?

    # 清理临时容器
    log_info "清理临时容器..." "$step"
    docker stop "$temp_container"
    docker rm "$temp_container"

    # 检查恢复状态
    if [ $RESTORE_STATUS -ne 0 ]; then
        handle_error 40 "PostgreSQL数据还原失败" "$step"
    fi

    log_success "在容器中完成PostgreSQL数据还原" "$step"
}

# 直接在主机上执行还原
restore_direct() {
    local step=$1
    local backup_name=$2

    log_info "直接在主机上执行WAL-G恢复..." "$step"

    # 确保WAL-G命令可用
    local walg_path=$(which wal-g)
    if [ -z "$walg_path" ]; then
        handle_error 7 "未找到wal-g命令，请先安装WAL-G" "$step"
    fi

    # 执行WAL-G恢复
    log_info "使用WAL-G恢复数据到 $POSTGRES_DATA_DIR..." "$step"
    ${walg_path} backup-fetch "$POSTGRES_DATA_DIR" "$backup_name"

    RESTORE_STATUS=$?
    if [ $RESTORE_STATUS -ne 0 ]; then
        handle_error 40 "PostgreSQL数据还原失败，退出代码: $RESTORE_STATUS" "$step"
    fi

    log_success "直接恢复PostgreSQL数据完成" "$step"
}

# 从增量备份还原
restore_incremental_backup() {
    local step=$1
    local backup_name=$2

    log_info "从备份 $backup_name 还原PostgreSQL数据..." "$step"

    # 根据直接恢复标志决定恢复模式
    if [ "$USE_DIRECT_RESTORE" = "true" ]; then
        # 直接在主机上执行恢复
        restore_direct "$step" "$backup_name"
    else
        # 在Docker容器中执行恢复
        if [ "$USE_DOCKER" = "true" ]; then
            restore_in_container "$step" "$backup_name"
        else
            # 本地PostgreSQL也使用直接恢复
            restore_direct "$step" "$backup_name"
        fi
    fi
}

# 设置权限
set_permissions() {
    local step=$1
    log_info "设置数据目录权限..." "$step"

    if [ "$USE_DOCKER" = "true" ]; then
        # Docker环境下设置正确的权限
        chmod 700 "$POSTGRES_DATA_DIR"
        chown -R 999:999 "$POSTGRES_DATA_DIR"  # Docker中postgres用户的UID/GID通常是999:999
    else
        # 本地PostgreSQL需要正确的权限
        chown -R postgres:postgres "$POSTGRES_DATA_DIR"
        chmod 700 "$POSTGRES_DATA_DIR"
    fi

    log_success "数据目录权限设置完成" "$step"
}

# 创建恢复信号文件
create_recovery_signal() {
    local step=$1
    log_info "创建恢复信号文件..." "$step"

    touch "$POSTGRES_DATA_DIR/recovery.signal"
    chmod 600 "$POSTGRES_DATA_DIR/recovery.signal"

    # 如果是本地PostgreSQL，需要设置正确的所有权
    if [ "$USE_DOCKER" != "true" ]; then
        chown postgres:postgres "$POSTGRES_DATA_DIR/recovery.signal"
    fi

    log_success "恢复信号文件已创建" "$step"
}

# 启动服务
start_services() {
    local step=$1
    log_info "启动服务..." "$step"

    # 启动PostgreSQL
    if [ "$USE_DOCKER" = "true" ]; then
        log_info "启动PostgreSQL容器..." "$step"
        docker start "$POSTGRES_CONTAINER_NAME"
        sleep 10

        # 检查容器是否正常运行
        if ! docker ps | grep -q "$POSTGRES_CONTAINER_NAME"; then
            handle_error 50 "PostgreSQL容器启动失败" "$step"
        fi
    else
        log_info "启动PostgreSQL服务..." "$step"
        systemctl start postgresql
        sleep 10

        if ! systemctl is-active postgresql &>/dev/null; then
            handle_error 51 "PostgreSQL服务启动失败" "$step"
        fi
    fi

    # 检查PostgreSQL连接
    check_postgres_available

    # 启动ThingsBoard
    log_info "启动ThingsBoard服务..." "$step"
    systemctl start thingsboard
    sleep 10

    if ! systemctl is-active thingsboard &>/dev/null; then
        log_warn "ThingsBoard服务启动失败，可能需要手动检查" "$step"
    else
        log_success "ThingsBoard服务启动成功" "$step"
    fi

    log_success "所有服务已启动" "$step"
}

# 发送通知
send_restore_notification() {
    local step=$1
    local backup_name=$2

    log_info "发送还原完成通知..." "$step"
    if [ "$ENABLE_NOTIFICATIONS" = "true" ]; then
        send_notification "ThingsBoard还原完成" "从备份 $backup_name 还原已完成"
        log_success "还原通知已发送" "$step"
    else
        log_info "通知功能未启用，跳过发送" "$step"
    fi
}

# ========== 4. 主程序入口 ==========

# 设置默认值
BACKUP_MODE=""
USE_DIRECT_RESTORE="false"  # 默认不使用直接恢复模式

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    key="$1"
    case $key in
        -h|--help)
            show_usage
            ;;
        -l|--list)
            list_backups
            ;;
        -f|--full)
            BACKUP_MODE="full"
            shift
            ;;
        -d|--diff)
            BACKUP_MODE="diff"
            shift
            ;;
        -D|--direct)
            USE_DIRECT_RESTORE="true"
            shift
            ;;
        *)
            BACKUP_NAME="$1"
            shift
            ;;
    esac
done

# 检查是否提供了备份名称
if [ -z "$BACKUP_NAME" ]; then
    echo "错误: 未指定备份名称"
    show_usage
fi

# 主要恢复流程
main() {
    log_info "================== ThingsBoard增量恢复操作开始 =================="
    if [ "$USE_DIRECT_RESTORE" = "true" ]; then
        log_info "使用直接恢复模式（在主机上执行WAL-G）"
    else
        log_info "使用标准恢复模式"
    fi

    # 步骤0: 检查环境
    check_environment

    # 步骤1: 检查备份位置
    if [ -z "$BACKUP_MODE" ]; then
        detect_backup_location
    fi

    # 如果指定了"latest"，获取最新备份
    step=$(next_step)
    if [ "$BACKUP_NAME" = "latest" ]; then
        get_latest_backup "$step"
    else
        log_info "使用指定的备份: $BACKUP_NAME" "$step"
    fi

    # 步骤2: 停止服务
    step=$(next_step)
    stop_services "$step"

    # 步骤3: 准备数据目录
    step=$(next_step)
    prepare_data_directory "$step"

    # 步骤4: 从增量备份还原
    step=$(next_step)
    restore_incremental_backup "$step" "$BACKUP_NAME"

    # 步骤5: 设置权限
    step=$(next_step)
    set_permissions "$step"

    # 步骤6: 创建恢复信号文件
    step=$(next_step)
    create_recovery_signal "$step"

    # 步骤7: 启动服务
    step=$(next_step)
    start_services "$step"

    # 步骤8: 完成恢复
    step=$(next_step)
    log_success "从备份 $BACKUP_NAME 还原已完成" "$step"
    log_info "您的ThingsBoard实例现在应该已恢复到备份时的状态" "$step"

    # 步骤9: 发送通知
    step=$(next_step)
    send_restore_notification "$step" "$BACKUP_NAME"

    log_info "================== ThingsBoard增量恢复操作完成 =================="
}

# 执行主恢复流程
main

exit 0
